package limiter;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.*;

/**
 * @author dushougudu
 */
public class TokenBucketLimiter {
    /**
     * 每秒处理数（放入令牌数量）
     */
    private long putTokenRate;

    /**
     * 最后刷新时间
     */
    private long refreshTime;

    /**
     * 令牌桶容量
     */
    private long capacity;

    /**
     * 当前桶内令牌数
     */
    private long currentToken = 0L;


    public TokenBucketLimiter(long capacity, long putTokenRate) {
        this.capacity = capacity;
        this.putTokenRate = putTokenRate;
    }

    /**
     * 漏桶算法
     *
     * @return
     */
    public synchronized boolean tokenBucketTryAcquire() {
        //获取系统当前时间
        long currentTime = System.currentTimeMillis();
        //生成的令牌 =(当前时间-上次刷新时间)* 放入令牌的速率
        long generateToken = (currentTime - refreshTime) / 1000 * putTokenRate;
        // 当前令牌数量 = 之前的桶内令牌数量+放入的令牌数量
        currentToken = Math.min(capacity, generateToken + currentToken);
        // 刷新时间
        refreshTime = currentTime;

        //桶里面还有令牌，请求正常处理
        if (currentToken > 0) {
            //令牌数量-1
            currentToken--;
            return true;
        }

        return false;
    }

    public static void main(String[] args) throws ExecutionException, InterruptedException {
        TokenBucketLimiter tokenBucketLimiter = new TokenBucketLimiter(5, 5);

        CyclicBarrier barrier = new CyclicBarrier(5);
        List<CompletableFuture<Boolean>> list = new ArrayList<>();
        ThreadPoolExecutor executor = new ThreadPoolExecutor(10, 10,
                0L, TimeUnit.MILLISECONDS,
                new LinkedBlockingQueue<Runnable>());
        for (int i = 0; i < 10; i++) {
            if (i == 5) {
                Thread.sleep(1000);
            }
            CompletableFuture<Boolean> future = CompletableFuture.supplyAsync(() -> {
                try {
                    barrier.await();
                } catch (InterruptedException e) {
                    e.printStackTrace();
                } catch (BrokenBarrierException e) {
                    e.printStackTrace();
                }
                boolean flag = tokenBucketLimiter.tokenBucketTryAcquire();
                return flag;
            }, executor);
            list.add(future);
        }

        for (CompletableFuture<Boolean> future : list) {
            System.out.println(future.get());
        }

        executor.shutdown();
    }


}
